from math import floor

import torch

from algorithms.algorithm import Algorithm
from algorithms.large.mocgde import MOCGDE
from individual.solution import Solution
from problems.problem import Problem
from utils.vectors import get_ref_vectors


class FRCGM(Algorithm):
    parameters = [
        {'label': 'beta', 'name': 'beta', 'type': 'number', 'description': 'A parameter within [0,1] for line search',
         'step': 0.01, 'defaultValue': 0.6},
        {'label': 'sigma', 'name': 'sigma', 'type': 'number',
         'description': 'A parameter within [0 0.5] for line search', 'step': 0.01, 'defaultValue': 0.4}
    ]

    def __init__(self, problem: Problem,
                 pop_size: int,
                 kwargs,
                 ):
        self.ref_vectors = get_ref_vectors(pop_size, problem.obj_dim)
        pop_size = self.ref_vectors.size(0)
        super().__init__(pop_size, problem, kwargs)
        self.population = problem.init_population(pop_size)
        self.k = 0
        self.g0 = [None] * pop_size
        self.d0 = [None] * pop_size

    def each_iteration(self):
        for i in range(self.pop_size):
            gk = MOCGDE.finite_difference(self.population[i], self.problem, self.ref_vectors[i])[0].squeeze(1)
            itern = self.k - (self.problem.var_dim + 1) * floor(self.k / (self.problem.var_dim + 1)) + 1

            if itern <= 1:
                dk = -gk
            else:
                betak = torch.dot(gk, gk) / torch.dot(self.g0[i], self.g0[i])
                dk = -gk + betak * self.d0[i]
                gd = torch.dot(gk, dk)
                if gd >= 0:
                    dk = -gk

            for m in range(21):
                dec = self.population[i].dec + self.beta ** m * dk
                dec = self.problem.repair_decision(dec)
                obj = self.problem.estimate_population(dec).obj.squeeze(0)
                if torch.dot(obj, self.ref_vectors[i]) <= torch.dot(self.population[i].obj, self.ref_vectors[i]) \
                        + self.sigma * self.beta ** m * torch.dot(gk, dk):
                    break
            self.population[i] = Solution.solution_build(dec, obj)
            self.g0[i] = gk
            self.d0[i] = dk
        self.k += 1
